import torch
from torch import nn
import math


def get_padding_mask(x,padding_idx):
    return (x == padding_idx).unsqueeze(1)

class Embedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.lut = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

class PositionalEncoding(nn.Module):
    #max_len表示句子的长度
    def __init__(self,d_model,max_len=5000,dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(max_len,d_model)
        pos = torch.arange(0,max_len).unsqueeze(1)
        div_term = torch.pow(10000,torch.arange(0,d_model,2).float() / d_model)
        #偶数列
        pe[:,0::2] = torch.sin(pos / div_term)
        #奇数列
        pe[:,1::2] = torch.cos(pos / div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
    def forward(self, x):
        # print(self.pe[:,:x.size(1)])
        x = x + self.pe[:,:x.size(1)]
        return self.dropout(x)


def attention(q,k,v,mask=None,dropout=None):
    #将k矩阵的最后一个维度值作为d_k
    d_k = k.size(-1)
    #将k矩阵的最后两个维度互换(转置),与q矩阵相乘,除以d_k开根号
    scores = torch.matmul(q,k.transpose(-2,-1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask,-1e9)
    p_attn = torch.softmax(scores,dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn,v)

class MultHeadedAttention(nn.Module):
    def __init__(self, d_model, n_head, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        # 必须能被整除
        assert d_model % n_head == 0
        self.d_token = d_model // n_head
        self.n_head = n_head

        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)

        self.linear = nn.Linear(d_model, d_model, bias=False)
        #layernorm
        self.norm = nn.LayerNorm(d_model)
    def forward(self, q, k, v,mask=None):
        res = q
        # 分头的主要目的是让模型能够同时关注输入序列的不同部分，从而捕捉更丰富的特征
        # q,k,v shape ==> batch_size,seq_len,d_mdoel
        n_batchs = q.size(0)
        q_seq_len = q.size(1)
        k_seq_len = k.size(1)
        v_seq_len = v.size(1)
        # (batch_size, seq_len, d_model) ===> (batch_size, seq_len, n_head, d_token)
        # 其中 n_head 是头的数量，d_token 是每个头的维度
        # 为了方便计算，需要将头的维度移到前面，变成 (batch_size, n_head, seq_len, d_token)
        # 矩阵运算（如 torch.matmul）通常对最后两个维度进行操作。
        # 将头的维度移到前面后，可以直接对 seq_len 和 d_token 维度进行矩阵乘法
        q = self.W_Q(q).view(n_batchs, q_seq_len, self.n_head, self.d_token).transpose(1, 2)
        k = self.W_K(k).view(n_batchs, k_seq_len, self.n_head, self.d_token).transpose(1, 2)
        v = self.W_V(v).view(n_batchs, v_seq_len, self.n_head, self.d_token).transpose(1, 2)

        # 掩码要升维
        if mask is not None:
            mask = mask.unsqueeze(1)
        # 计算注意力
        att = attention(q, k, v,mask,self.dropout)
        # 拼接
        # contiguous()确保张量在内存中是连续存储的，提高计算效率；
        # (batch_size, n_head, seq_len, d_token)=>(batch_size, seq_len, n_head, d_token)=>(batch_size, seq_len, d_model)
        concat = att.transpose(1, 2).contiguous().reshape(n_batchs, -1, self.n_head * self.d_token)
        output = self.linear(concat)
        return self.norm(res + output)

class FeedForwardNet(nn.Module):
    def __init__(self,d_model,d_ff,dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.w1 = nn.Linear(d_model,d_ff)
        self.w2 = nn.Linear(d_ff,d_model)
        self.relu = nn.ReLU()
        self.norm = nn.LayerNorm(d_model)

    def forward(self,x):
        res = x
        x = self.w1(x)
        x = self.relu(x)
        x = self.w2(x)
        x = self.dropout(x)
        return self.norm(x+res)

class EncoderLayer(nn.Module):
    def __init__(self,d_model,n_head,d_ff,dropout=0.1):
        super().__init__()
        self.mha = MultHeadedAttention(d_model,n_head,dropout)
        self.ffn = FeedForwardNet(d_model,d_ff,dropout)
    def forward(self,x,mask=None):
        output = self.mha(x,x,x,mask)
        return self.ffn(output)

class Encoder(nn.Module):
    def __init__(self,vocal_size,d_model,n_head,d_ff,N=6,dropout=0.1):
        super().__init__()
        self.emb = Embedding(vocal_size,d_model)
        self.pos = PositionalEncoding(d_model=d_model,dropout=dropout)
        self.layers = nn.ModuleList([EncoderLayer(d_model,n_head,d_ff,dropout) for _ in range(N)])

    def forward(self,x,mask=None):
        x = self.emb(x)
        x = self.pos(x)
        for layer in self.layers:
            x = layer(x,mask)
        return x

# inputs = torch.tensor([
#     [1,2,3],
#     [4,5,0]
# ])
# mask = get_padding_mask(inputs,0)
# enc = Encoder(10,8,2,32)
# enc_out = enc(inputs,mask)
# print(enc_out.shape)
# print(enc_out)